/*
 * Copyright 2010, 2011 Red Hat, Inc. and/or its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package core.impl.p7;

import api.kiecontainer.Environment;
import api.kiecontainer.KieBase;
import api.kiecontainer.KieSession;
import api.kiecontainer.KieSessionConfiguration;
import api.kiecontainer.SessionClock;
import api.kieservices.MarshallingConfiguration;
import api.kieservices.ObjectMarshallingStrategyStore;
import core.api.a1.InternalKnowledgeBase;
import core.api.a1.InternalWorkingMemory;
import core.api.a1.KieSessionInitializer;
import core.api.a3.InternalMarshaller;
import core.api.a3.TimersInputMarshaller;
import core.impl.p1.RuleBaseConfiguration;
import core.impl.p1.SlidingTimeWindow;
import core.impl.p1.StatefulKnowledgeSessionImpl;
import core.impl.p2.MarshallerReaderContext;
import core.impl.p2.MarshallerWriteContext;
import core.impl.p2.ObjectTypeNode;
import core.impl.p2.PhreakTimerNode;
import core.impl.p2.ProtobufInputMarshaller;
import core.impl.p2.ProtobufMessages;
import core.impl.p2.ProtobufOutputMarshaller;
import core.impl.p2.SessionConfiguration;
import core.impl.p6.KnowledgeBaseImpl;
import internal.api.a1.StatefulKnowledgeSession;
import internal.impl.p1.KnowledgeBaseFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.Map;

/**
 * A Marshaller implementation that uses ProtoBuf as the marshalling
 * framework in order to support backward compatibility with
 * marshalled sessions
 * 
 */
public class ProtobufMarshaller
        implements
        InternalMarshaller {

    private KieSessionInitializer initializer;

    public KieSessionInitializer getInitializer() {
        return initializer;
    }

    public void setInitializer( KieSessionInitializer initializer ) {
        this.initializer = initializer;
    }

    public static final Map<Integer, TimersInputMarshaller> TIMER_READERS = new HashMap<Integer, TimersInputMarshaller>();
    static {
        TIMER_READERS.put( ProtobufMessages.Timers.TimerType.BEHAVIOR_VALUE, new SlidingTimeWindow.BehaviorJobContextTimerInputMarshaller() );
        TIMER_READERS.put( ProtobufMessages.Timers.TimerType.EXPIRE_VALUE, new ObjectTypeNode.ExpireJobContextTimerInputMarshaller() );
        TIMER_READERS.put( ProtobufMessages.Timers.TimerType.TIMER_NODE_VALUE, new PhreakTimerNode.TimerNodeTimerInputMarshaller() );
    }
    
    KieBase kbase;
    RuleBaseConfiguration ruleBaseConfig;
    MarshallingConfiguration marshallingConfig;
    ObjectMarshallingStrategyStore strategyStore;

    public ProtobufMarshaller(KieBase kbase,
                              MarshallingConfiguration marshallingConfig) {
        this.kbase = kbase;
        this.ruleBaseConfig = RuleBaseConfiguration.getDefaultInstance();
        this.marshallingConfig = marshallingConfig;
        this.strategyStore = this.marshallingConfig.getObjectMarshallingStrategyStore();
    }

    public StatefulKnowledgeSession unmarshall(final InputStream stream) throws IOException,
                                                                        ClassNotFoundException {
        return unmarshall( stream, null, null );
    }

    public StatefulKnowledgeSession unmarshall(final InputStream stream,
                                               KieSessionConfiguration config,
                                               Environment environment) throws IOException,
                                                                       ClassNotFoundException {
        if ( config == null ) {
            config = KnowledgeBaseFactory.newKnowledgeSessionConfiguration();
        }

        if ( environment == null ) {
            environment = KnowledgeBaseFactory.newEnvironment();
        }

        MarshallerReaderContext context = new MarshallerReaderContext( stream,
                                                                       (KnowledgeBaseImpl) kbase,
                                                                       RuleBaseNodes.getNodeMap( (KnowledgeBaseImpl) kbase ),
                                                                       this.strategyStore,
                                                                       TIMER_READERS,
                                                                       this.marshallingConfig.isMarshallProcessInstances(),
                                                                       this.marshallingConfig.isMarshallWorkItems(),
                                                                       environment );

        int id = ((KnowledgeBaseImpl) this.kbase).nextWorkingMemoryCounter();
        RuleBaseConfiguration conf = ((KnowledgeBaseImpl) this.kbase).getConfiguration();

        StatefulKnowledgeSessionImpl session = ProtobufInputMarshaller.readSession(context,
                id,
                environment,
                (SessionConfiguration) config,
                initializer);
        context.close();
        if ( ((SessionConfiguration) config).isKeepReference() ) {
            ((KnowledgeBaseImpl) this.kbase).addStatefulSession(session);
        }
        return session;

    }

    public void unmarshall(final InputStream stream,
                           final KieSession ksession) throws IOException,
                                                                   ClassNotFoundException {
        MarshallerReaderContext context = new MarshallerReaderContext( stream,
                                                                       (KnowledgeBaseImpl) kbase,
                                                                       RuleBaseNodes.getNodeMap( (KnowledgeBaseImpl) kbase ),
                                                                       this.strategyStore,
                                                                       TIMER_READERS,
                                                                       this.marshallingConfig.isMarshallProcessInstances(),
                                                                       marshallingConfig.isMarshallWorkItems(),
                                                                       ksession.getEnvironment() );

        ProtobufInputMarshaller.readSession((StatefulKnowledgeSessionImpl) ksession,
                                            context);
        context.close();
    }

    public void marshall(final OutputStream stream,
                         final KieSession ksession) throws IOException {
        marshall( stream, ksession, ksession.<SessionClock> getSessionClock().getCurrentTime() );
    }

    public void marshall(final OutputStream stream,
                         final KieSession ksession,
                         final long clockTime) throws IOException {
        ((InternalWorkingMemory) ksession).flushPropagations();
        MarshallerWriteContext context = new MarshallerWriteContext( stream,
                                                                     (InternalKnowledgeBase) kbase,
                                                                     (InternalWorkingMemory) ksession,
                                                                     RuleBaseNodes.getNodeMap( (InternalKnowledgeBase) kbase ),
                                                                     this.strategyStore,
                                                                     this.marshallingConfig.isMarshallProcessInstances(),
                                                                     this.marshallingConfig.isMarshallWorkItems(),
                                                                     ksession.getEnvironment() );
        context.clockTime = clockTime;
        ProtobufOutputMarshaller.writeSession(context);
        context.close();
    }

    public MarshallingConfiguration getMarshallingConfiguration() {
        return marshallingConfig;
    }

}
